/*
 * License: Copyright (c) Huawei Technologies Co., Ltd. 2012-2019. All rights reserved.
 * Description: video encode module
 * Date: 2020-01-07 09:50:36
 * LastEditTime : 2020-02-19 16:28:51
 */

#include <queue>
#include "hiaiengine/log.h"
#include "venc_engine.h"
#include <sys/time.h>
#include "hiaiengine/ai_memory.h"

const int KEY_FRAME_INTERVAL = 250;
const int VENC_SCALE_LOWER_BOUND = 128;
const int VENC_SCALE_UPPER_BOUND = 1920;
const int TIMES_SECOND_MICROSECOND = 1000000;

uint64_t VencEngine::GetCurentTimeStamp()
{
    gettimeofday(&currentTimeval, NULL);
    return currentTimeval.tv_sec * TIMES_SECOND_MICROSECOND + currentTimeval.tv_usec;
}

void VencEngine::VencCallBack(struct VencOutMsg* vencOutMsg, void* userData)
{
    HIAI_StatusT ret;
    if (userData == nullptr) {
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, " input parameter of VencCallback error!");
        return;
    }

    auto enginePtr = (VencEngine*)userData;
    if (!vencOutMsg) {
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, " VencCallback received null error");
        return;
    }

    uint8_t* outputBuffer = nullptr;
    outputBuffer = new uint8_t[vencOutMsg->outputDataSize];
    if (outputBuffer == nullptr) {
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "VENC new memory error! ");
        return;
    }

    ret = memcpy_s(outputBuffer, vencOutMsg->outputDataSize, vencOutMsg->outputData, vencOutMsg->outputDataSize);
    if (ret != 0) {
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "VENC memcpy error! ErrCode: %d", ret);
        delete[] outputBuffer;
        outputBuffer = nullptr;
        return;
    }

    std::shared_ptr<TransData> encodedDataCache = std::make_shared<TransData>();
    encodedDataCache->buf.reset(outputBuffer, [](uint8_t* p) {delete[] p;});
    encodedDataCache->len = vencOutMsg->outputDataSize;
    encodedDataCache->frameId = (enginePtr->frameIndex)++;
    enginePtr->encodedQueue.push(encodedDataCache);
    return;
}

VencEngine::~VencEngine()
{
    if (vencHandle != -1) {
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "Destroy Venc api!");
        DestroyVenc(vencHandle);
        vencHandle = -1;
    }
}

HIAI_StatusT VencEngine::VencInit(std::shared_ptr<TransData> transData)
{
    uint32_t width = transData->alignWidth;
    uint32_t height = transData->alignHeight;
    int videoFormat = transData->videoFormat;

    if (width < VENC_SCALE_LOWER_BOUND || width > VENC_SCALE_UPPER_BOUND || width % 2 != 0) {
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "width of Venc is invalid, please check! ");
        return HIAI_ERROR;
    }

    if (height < VENC_SCALE_LOWER_BOUND || height > VENC_SCALE_UPPER_BOUND || height % 2 != 0) {
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "height of Venc is invalid, please check! ");
        return HIAI_ERROR;
    }

    struct VencConfigV2 vencConfig;
    vencConfig.width = width;
    vencConfig.height = height;
    vencConfig.codingType = videoFormat ? 0 : 2; // 0: h265 main, 2: h264 main
    vencConfig.yuvStoreType = 0; // YUV420SP
    vencConfig.keyFrameInterval = KEY_FRAME_INTERVAL; // the interval between I frames
    vencConfig.vencOutMsgCallBack = VencEngine::VencCallBack;
    vencConfig.userData = (void*)this;
    vencConfig.controlRateType = transData->vencCtlType;
    vencConfig.bitRate = transData->vencBitRate;
    vencConfig.inSrcFrameRate = transData->vencFrameRate;
    vencHandle = CreateVencV2(&vencConfig);
    if (vencHandle == -1) {
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "CreateVenc fail!");
        return HIAI_ERROR;
    }

    return HIAI_OK;
}

HIAI_IMPL_ENGINE_PROCESS("VencEngine", VencEngine, INPUT_SIZE)
{
    std::shared_ptr<TransData> transData = std::static_pointer_cast<TransData>(arg0);
    if (transData == nullptr) {
        return HIAI_ERROR;
    }

    if (transData->eos) {
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "[VencEngine] Get the end of the video!");
        if (hiai::Engine::SendData(0, "TransData", std::static_pointer_cast<void>(transData)) != HIAI_OK) {
            HIAI_ENGINE_LOG(HIAI_IDE_ERROR, " SendData of VencEngine error!");
            return HIAI_ERROR;
        }
        return HIAI_OK;
    }

    if (vencHandle == -1) {
        if (HIAI_OK != VencInit(transData)) {
            return HIAI_ERROR;
        }
        initialTimeStamp = GetCurentTimeStamp();
    }

    if (EncodeYuv2Video(transData) != HIAI_OK) {
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "[VencEngine] Encode error!");
        return HIAI_ERROR;
    }

    while (!encodedQueue.empty()) {
        std::shared_ptr<TransData> encodedTemp = encodedQueue.front();
        encodedQueue.pop();

        // calculate the time stamp
        std::shared_ptr<TransData> rawDataOutput = std::make_shared<TransData>();
        rawDataOutput->curEncodeFps = (encodedTemp->frameId) /
            (1.0 * (GetCurentTimeStamp() - initialTimeStamp) / TIMES_SECOND_MICROSECOND);

        // construct data transmission struct
        rawDataOutput->buf = encodedTemp->buf;
        rawDataOutput->len = encodedTemp->len;
        rawDataOutput->alignWidth = transData->alignWidth;
        rawDataOutput->alignHeight = transData->alignHeight;
        rawDataOutput->frameId = encodedTemp->frameId;
        rawDataOutput->graphId = transData->graphId;
        rawDataOutput->deviceId = transData->deviceId;
        rawDataOutput->channelId = transData->channelId;
        rawDataOutput->outputFormat = transData->outputFormat;
        rawDataOutput->videoFormat = transData->videoFormat;
        rawDataOutput->curDecodeFps = transData->curDecodeFps;
        rawDataOutput->streamPullFps = transData->streamPullFps;

        if (hiai::Engine::SendData(0, "TransData", std::static_pointer_cast<void>(rawDataOutput)) !=
            HIAI_OK) {
            HIAI_ENGINE_LOG(HIAI_IDE_ERROR, " SendData of VencEngine error!");
            return HIAI_ERROR;
        }
    }
    return HIAI_OK;
}

HIAI_StatusT VencEngine::EncodeYuv2Video(std::shared_ptr<TransData> transData)
{
    struct VencInMsg vencInMsg;
    vencInMsg.inputData = transData->buf.get();
    vencInMsg.inputDataSize = transData->len;
    vencInMsg.keyFrameInterval = KEY_FRAME_INTERVAL;
    vencInMsg.forceIFrame = 0;
    vencInMsg.eos = 0;

    if (RunVenc(vencHandle, &vencInMsg) == -1) {
        HIAI_ENGINE_LOG(HIAI_IDE_ERROR, "call video encode fail");
        return HIAI_ERROR;
    }
    return HIAI_OK;
}
